import threading
import queue
import heapq


class Router(threading.Thread):
    def __init__(self, router_ip, router_mac, log_queue):
        super().__init__()
        self.router_id = router_ip  # 路由器名称 模拟中使用id代替ip地址
        self.router_mac = router_mac # 本机mac地址
        self.log_queue = log_queue  # 本机操作日志
        self.router_table = {}  # 本机路由表
        self.arp_table = {}  # 本机ARP表
        self.link_state_database = {}  # 链路状态数据库，保存路由器之间的链路状态信息(名称和距离)

    def add_arp(self, arp_ip, arp_mac):  # 添加arp表记录
        self.arp_table[arp_ip] = arp_mac

    def construct_path(self, destination_router_id, predecessors):
        path = []
        current_router_id = destination_router_id
        while current_router_id is not None:
            path.insert(0, current_router_id)
            current_router_id = predecessors[current_router_id]
        return path

    def calculate_shortest_paths(self): # 计算最短路径
        # 初始化距离和前驱节点
        distances = {router_id: float("inf") for router_id in self.link_state_database}
        distances[self.router_id] = 0
        predecessors = {router_id: None for router_id in self.link_state_database}

        # 使用Dijkstra算法计算最短路径
        queue = [(0, self.router_id)]  # 优先队列，用于选择下一个要处理的节点
        while queue:
            curr_distance, curr_router = heapq.heappop(queue)
            if curr_distance > distances[curr_router]:
                continue

            # 遍历当前节点的邻居
            if curr_router in self.link_state_database:
                for neighbor_router, cost in self.link_state_database[curr_router].items():
                    new_distance = curr_distance + cost
                    if new_distance < distances[neighbor_router]:
                        distances[neighbor_router] = new_distance
                        predecessors[neighbor_router] = curr_router
                        heapq.heappush(queue, (new_distance, neighbor_router))

        # 构建最短路径字典
        shortest_paths = {}
        for destination_router_id in self.link_state_database:
            path = self.construct_path(destination_router_id, predecessors)
            shortest_paths[destination_router_id] = path

        return shortest_paths

    def build_routing_table(self): # 完善路由表
        # 使用Dijkstra算法计算最短路径
        shortest_paths = self.calculate_shortest_paths()

        # 构建路由表
        for destination_router_id, path in shortest_paths.items():
            if len(path) >= 2:
                next_hop = path[1]  # 下一跳路由器
                self.router_table[destination_router_id] = next_hop
            else:
                # 没有下一跳，该路由器是目标路由器本身
                self.router_table[destination_router_id] = destination_router_id

    def add_linkstate(self, nearby_router_id, cost):
        if self.router_id not in self.link_state_database:
            self.link_state_database[self.router_id] = {}
        self.link_state_database[self.router_id][nearby_router_id] = cost

        if nearby_router_id not in self.link_state_database:
            self.link_state_database[nearby_router_id] = {}
        self.link_state_database[nearby_router_id][self.router_id] = cost

    def send_link_state_update(self): # 发送链路状态更新报文给相邻路由器
        for neighbor_router_id in self.link_state_database[self.router_id]:
            link_state_info = self.link_state_database[self.router_id]
            self.log_queue.put((self.router_id, neighbor_router_id, link_state_info))

    def receive_link_state_update(self, source_router_id, link_state_info):
        self.link_state_database[source_router_id] = link_state_info

    def process_frame(self, frame, source_mac):

        # 解析以太网帧，提取数据部分
        data = frame["data"]
        if isinstance(data, dict):
            # 处理ARP分组
            if "opcode" in data:
                self.process_arp_packet(data, source_mac)
            # 处理IPv4分组
            elif "source_ip" in data and "destination_ip" in data:
                self.process_ipv4_packet(data)
            # 处理ICMP分组
            elif "ICMP" in data:
                self.process_icmp_packet(data)
            # 处理其他类型的分组
            else:
                log_message = f"No action taken for unknown packet on {self.router_id}"
                print(log_message)
        else:
            log_message = f"No action taken for unknown packet on {self.router_id}"
            print(log_message)

    def process_arp_packet(self, packet, source_mac):
        # 解析ARP分组
        source_ip = packet["source_ip"]
        target_ip = packet["target_ip"]

        # 处理ARP请求
        if packet["opcode"] == "request": # ARP请求
            # 判断目标IP地址是否为当前路由器的接口IP地址
            if target_ip in self.arp_table:
                # 构造ARP响应
                arp_response = {
                    "opcode": "reply",
                    "source_ip": target_ip,
                    "source_mac": self.arp_table.get(target_ip),
                    "target_ip": source_ip,
                }

                # 构造以太网帧
                ethernet_frame = {
                    "destination_mac": source_mac,
                    "source_mac": self.router_mac,
                    "data": arp_response
                }

                # 发送以太网帧
                self.log_queue.put(ethernet_frame)
                return

        elif packet["opcode"] == "reply":
            self.add_arp(packet["source_ip"], packet["source_mac"])
            if target_ip == self.router_id:
                log_message = f"{self.router_id} get arp_reply"
                print(log_message)
            else:
                next_hop = self.router_table.get(target_ip)
                if next_hop:
                    ethernet_frame = {
                        "source_mac": self.router_id,
                        "destination_mac": self.arp_table.get(next_hop),
                        "data": packet
                    }
                    self.log_queue.put(ethernet_frame)
            return

    def process_ipv4_packet(self, packet):
        destination_ip = packet["destination_ip"]

        if destination_ip == self.router_id:
            receives=packet["data"]
            log_message = f"{self.router_id} got IPv4 packet: \n{receives}"
            print(log_message)
            return

        # 进行IP路由选择
        next_hop = self.router_table.get(destination_ip)
        if next_hop:
            ethernet_frame = {
                "source_mac": self.router_id,
                "destination_mac": self.arp_table.get(next_hop),
                "data": packet
            }
            self.log_queue.put(ethernet_frame)

    def process_icmp_packet(self, packet):
        source_ip = packet["source_ip"]
        destination_ip = packet["destination_ip"]
        icmp_type = packet["ICMP"]

    def run(self):  # 路由器实际运行
        self.add_arp(self.router_id,self.router_mac)
        self.build_routing_table()
        while True:
            # 处理接收到的消息
            try:
                frame = self.log_queue.get(timeout=1)
                if frame["destination_mac"] == self.router_mac:
                    # 收到目标为自己的消息，处理数据包
                    self.process_frame(frame, frame["source_mac"])
                elif frame["destination_mac"] == '000':
                    self.process_frame(frame, frame["source_mac"])
                else:
                    self.log_queue.put(frame)
            except queue.Empty:
                # 队列为空，继续下一轮循环
                pass
            except TypeError:
                pass


class EthernetSwitch:
    def __init__(self, switch_name, router, log_queue):
        self.switch_name = switch_name
        self.router = router  # 连接的路由器
        self.mac_table = {}  # 交换表
        self.log_queue = log_queue

    def learn_mac_address(self, src_mac, in_port):
        self.mac_table[src_mac] = in_port

    def forward_frame(self, frame, in_port):
        dst_mac = frame['destination_mac']
        print(f"Switch: {self.switch_name} received frame from inport: {in_port}, frame: \n{frame}")
        if (dst_mac in self.mac_table) or (dst_mac == self.router.router_mac):
            out_port = self.mac_table[dst_mac]
            print(f"Switch {self.switch_name} forwarding frame from port {in_port} to port {out_port}")
            self.log_queue.put(frame)
        else:
            print("Destination MAC address not found in the MAC table. Broadcasting frame to all ports.")
            for macs, in_ports in self.mac_table.items():
                frame['destination_mac'] = macs
                self.log_queue.put(frame)
            frame['destination_mac'] = self.router.router_mac
            self.log_queue.put(frame)

    def process_frame(self, frame, in_port):
        src_mac = frame['source_mac']
        dst_mac = frame['destination_mac']

        self.learn_mac_address(src_mac, in_port)
        self.forward_frame(frame, in_port)


class Host(threading.Thread):
    def __init__(self, host_id, host_mac, switch, switch_inport, log_queue):
        super().__init__()
        self.host_id = host_id
        self.switch = switch
        self.switch_inport = switch_inport  # 交换机端口号
        self.host_mac = host_mac
        self.arp_table= {}
        self.log_queue = log_queue

    def add_arp(self, arp_ip, arp_mac):  # 添加arp表记录
        self.arp_table[arp_ip] = arp_mac

    def ask_mac(self, destination_ip):
        self.log_queue.put({'source_mac': self.host_mac, 'destination_mac': '000', 'data':
    {'source_ip': self.host_id, 'target_ip': destination_ip, 'opcode': 'request'}})

    def construct_ipv4(self, source_ip, destination_ip, data):
        ipv4_packet={
            'source_ip': source_ip,
            'destination_ip': destination_ip,
            'data': data
        }
        print(f"Host：{self.host_id} have constructed ipv4: {data}")
        destination_mac=self.arp_table.get(destination_ip)
        if destination_mac:
            self.send_frame(destination_mac, ipv4_packet)
        else:
            self.ask_mac(destination_ip)

    def send_frame(self, destination_mac, packet):
        # 构造以太网帧
        ethernet_frame = {
            "source_mac": self.host_mac,
            "destination_mac": destination_mac,
            "data": packet
        }
        self.switch.process_frame(ethernet_frame, self.switch_inport)

    def process_frame(self, frame, source_mac):
        log_message = f"{self.host_id} received Ethernet frame: \n{frame}"
        print(log_message)

        # 解析以太网帧，提取数据部分
        data = frame["data"]
        if isinstance(data, dict):
            # 处理ARP分组
            if "opcode" in data:
                self.process_arp_packet(data, source_mac)
            # 处理IPv4分组
            elif "source_ip" in data and "destination_ip" in data:
                self.process_ipv4_packet(data)
            # 处理ICMP分组
            elif "ICMP" in data:
                self.process_icmp_packet(data)
            # 处理其他类型的分组
            else:
                log_message = f"No action taken for unknown packet on {self.host_id}"
                print(log_message)
        else:
            log_message = f"No action taken for unknown packet on {self.host_id}"
            print(log_message)

    def process_arp_packet(self, packet, source_mac):
        # 解析ARP分组
        source_ip = packet["source_ip"]
        target_ip = packet["target_ip"]

        # 记录日志
        log_message = f"{self.host_id} received ARP packet: \n{packet}"
        print(log_message)

        # 处理ARP请求
        if packet["opcode"] == "request": # ARP请求
            # 判断目标IP地址是否为当前路由器的接口IP地址
            if target_ip in self.arp_table:
                # 构造ARP响应
                arp_response = {
                    "opcode": "reply",
                    "source_ip": target_ip,
                    "source_mac": self.arp_table.get(target_ip),
                    "target_ip": source_ip,
                }

                # 构造以太网帧
                self.send_frame(source_mac, arp_response)
                return

        elif packet["opcode"] == "reply":
            self.add_arp(packet["source_ip"], packet["source_mac"])
            return

    def process_ipv4_packet(self, packet):
        destination_ip = packet["destination_ip"]

        # 记录日志
        log_message = f"{self.host_id} received IPv4 packet: \n{packet}"
        print(log_message)

        if destination_ip == self.host_id:
            receives=packet["data"]
            log_message = f"{self.host_id} got IPv4 packet: \n{receives}\n"
            print(log_message)
            return

    def process_icmp_packet(self, packet):
        source_ip = packet["source_ip"]
        destination_ip = packet["destination_ip"]
        icmp_type = packet["ICMP"]

    def run(self):
        while True:
            try:
                frame = self.log_queue.get(timeout=1)
                if frame["destination_mac"] == self.host_mac:
                    # 收到目标为自己的消息，处理数据包
                    self.process_frame(frame, frame["source_mac"])
                elif frame["destination_mac"] == '000':
                    self.process_frame(frame, frame["source_mac"])
                else:
                    self.log_queue.put(frame)
            except queue.Empty:
                # 队列为空，继续下一轮循环
                pass
            except TypeError:
                pass
            # 处理接收到的消息


log_queue = queue.Queue()

# 创建路由器实例
router1 = Router("R1", "00:11:22:33:44:51", log_queue)
router2 = Router("R2", "00:11:22:33:44:52", log_queue)
router3 = Router("R3", "00:11:22:33:44:53", log_queue)

# 创建交换机实例
switch1 = EthernetSwitch("S1", router1, log_queue)
switch2 = EthernetSwitch("S2", router3, log_queue)

# 创建主机实例
host1 = Host("H1", "00:11:22:33:45:51", switch1, "A", log_queue)
host2 = Host("H2", "00:11:22:33:45:52", switch1, "B", log_queue)
host3 = Host("H3", "00:11:22:33:45:53", switch2, "A", log_queue)
host4 = Host("H4", "00:11:22:33:45:54", switch2, "B", log_queue)

# 添加链路状态信息
router1.add_linkstate("R2", 2)
router1.add_linkstate("R3", 5)
router2.add_linkstate("R1", 2)
router2.add_linkstate("R3", 1)
router3.add_linkstate("R1", 5)
router3.add_linkstate("R2", 1)

# 发送链路状态更新报文
router1.send_link_state_update()
router2.send_link_state_update()
router3.send_link_state_update()

# 启动路由器线程
router1.start()
router2.start()
router3.start()
# 启动主机线程
host1.start()
host2.start()
host3.start()
host4.start()

